from torch import linspace,randn_like,sqrt
from torch.nn.functional import mse_loss

def linear_beta_schedule(timesteps):
    beta_start = 0.0001
    beta_end = 0.02
    return linspace(beta_start, beta_end, timesteps)

def extract(a, t, x_shape):
    batch_size = t.shape[0]
    out = a.gather(-1, t.cpu())
    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)

def q_sample(x_start, t, alphas_cumprod, noise=None):
    if noise is None:
        noise = randn_like(x_start)
    sqrt_alphas_cumprod = sqrt(alphas_cumprod)
    sqrt_one_minus_alphas_cumprod = sqrt(1. - alphas_cumprod)
    sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
    sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t, x_start.shape)
    
    return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise

def p_losses(denoise_model, x_start, t, alphas_cumprod, noise=None):
    if noise is None:
        noise = randn_like(x_start)
    x_noisy = q_sample(x_start, t, alphas_cumprod, noise)
    predicted_noise = denoise_model(x_noisy, t)
    loss = mse_loss(noise, predicted_noise)
    return loss